import traceback
import re

from model_utils import load_model


ERROR_MSG = "Error, please retry"


def chat(model, tokenizer, processor, img, msgs, ctx, params=None, vision_hidden_states=None):
    default_params = {"num_beams":3, "repetition_penalty": 1.2, "max_new_tokens": 1024}
    if params is None:
        params = default_params
    if img is None:
        return -1, "Error, invalid image, please upload a new image", None, None
    try:
        image = img.convert('RGB')
        answer = model.chat(
            image=image, 
            msgs=msgs,
            tokenizer=tokenizer,
            **params
        )
        res = re.sub(r'(<box>.*</box>)', '', answer)
        res = res.replace('<ref>', '')
        res = res.replace('</ref>', '')
        res = res.replace('<box>', '')
        answer = res.replace('</box>', '')
        return -1, answer, None, None
    except Exception as err:
        print(err)
        traceback.print_exc()
        return -1, ERROR_MSG, None, None


def MiniCPM_respond(models_dir, model_name, _question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature):
    model, tokenizer, processor = load_model(models_dir, model_name)
    # 检查_app_cfg字典中是否存在'ctx'键，如果不存在则返回提示信息
    if _app_cfg.get('ctx', None) is None:
        # 将用户问题和提示信息添加到_chat_bot列表中
        _chat_bot.append((_question, 'Please upload an image to start'))
        # 返回空字符串、更新后的_chat_bot列表和未修改的_app_cfg字典
        return '', _chat_bot, _app_cfg

    # 复制_app_cfg字典中的'ctx'键对应的值，以避免直接修改原字典
    _context = _app_cfg['ctx'].copy()
    # 如果_context不为空，则将用户的问题添加到_context中
    if _context:
        _context.append({"role": "user", "content": _question})
    # 如果_context为空，则初始化_context并添加用户的问题
    else:
        _context = [{"role": "user", "content": _question}] 
    # 打印用户的问题
    print('<User>:', _question)

    # 根据params_form的值设置不同的参数
    if params_form == 'Beam Search':
        # 设置Beam Search算法的参数
        params = {
            'sampling': False,  # 不使用采样
            'num_beams': num_beams,  # Beam Search的束宽
            'repetition_penalty': repetition_penalty,  # 重复惩罚
            "max_new_tokens": 896   # 最大生成的新token数
        }
    else:
        # 设置其他生成算法的参数
        params = {
            'sampling': True,  # 使用采样
            'top_p': top_p,
            'top_k': top_k,  # 顶部k采样参数
            'temperature': temperature,  # 温度参数
            'repetition_penalty': repetition_penalty_2,  # 重复惩罚
            "max_new_tokens": 896   # 最大生成的新token数
        }
    # 调用chat函数进行对话生成，传入图片、上下文和参数
    code, _answer, _, sts = chat(model, tokenizer, processor, _app_cfg['img'], _context, None, params)
    print('<Assistant>:', _answer)

    # 将助手（assistant）的角色和回答内容添加到上下文（_context）列表中
    _context.append({"role": "assistant", "content": _answer}) 
    # 将问题和回答以元组形式添加到聊天机器人（_chat_bot）列表中
    _chat_bot.append((_question, _answer))
    # 检查代码（code）是否为0，如果是0，则更新应用配置（_app_cfg）
    if code == 0:
        # 将当前上下文（_context）存储到应用配置的'ctx'键中
        _app_cfg['ctx']=_context
        # 将状态（sts）存储到应用配置的'sts'键中
        _app_cfg['sts']=sts
    # 返回一个空字符串，聊天机器人列表和应用配置
    return '', _chat_bot, _app_cfg